
# -*- coding: utf-8 -*-


import unittest
import sys
import os
import time
import warnings
import logging

# ================================================================================
# 压测RPC需要的类
# 生成报告
# jmeter -n -t "HTTP Request.jmx" -l result.jtl -e -o /Users/yaozhang/Desktop/test-report/report

import grpc
from locust import events, User, task
from locust.exception import LocustError
from locust.user.task import LOCUST_STATE_STOPPING
import gevent
import time
import grpc.experimental.gevent as grpc_gevent

grpc_gevent.init_gevent()


log_fmt = "[%(levelname)s]%(asctime)s line %(lineno)d :\n%(message)s"
c_fmt = "[%(levelname)s]%(asctime)s %(filename)s.%(funcName)s():line %(lineno)d :\n%(message)s"
date_format = "%Y-%m-%d %H:%M:%S %a"

# 设置控制台输出level
logging.basicConfig(level=logging.INFO,
                    format=c_fmt,
                    datefmt=date_format
                    )


host = "127.0.0.1"
port = "9898"
# ================================================================================


# 添加执行环境
file_exec_path = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
sys.path.insert(0, file_exec_path)
from framework_strategy.gopy.trader.engine import EmailEngine, MainEngine, EventEngine  # NOQA
from framework_strategy.gopy.base.api import API, Event  # NOQA
from framework_strategy.gopy.gateways.object import *  # NOQA
from framework_strategy.gopy.gateways.xt.xt_api import XTGateway  # NOQA
from framework_strategy.manager import *  # NOQA
from framework_strategy.gopy.gateways.proto import message_proto_pb2, message_proto_pb2_grpc  # NOQA
from concurrent import futures

def init_engine():
    event_engine = EventEngine()
    xt_gateway = XTGateway(event_engine, "BINANCE")
    event_engine.start()
    # 获取服务器时间同步获取
    api = API(event_engine)
    return event_engine, api




def run():
    """test grpc server demo """
    with grpc.insecure_channel(host + ":" + port) as channel:

        stub = message_proto_pb2_grpc.EchoProtoServerStub(channel)
        method = "get_account_info"
        response = stub.request(
            message_proto_pb2.GateWayRequest(method=method))

    print("[ddddddddddddddddddddddddddddddddddd]", response.data)
    print(response.data)  # str


class GrpcClient:
    def __init__(self, environment, stub):
        self.env = environment
        self._stub_class = stub.__class__
        self._stub = stub

    def __getattr__(self, name):
        func = self._stub_class.__getattribute__(self._stub, name)

        def wrapper(*args, **kwargs):
            request_meta = {
                "request_type": "grpc",
                "name": name,
                "start_time": time.time(),
                "response_length": 0,
                "exception": None,
                "context": None,
                "response": None,
            }
            start_perf_counter = time.perf_counter()
            try:
                request_meta["response"] = func(*args, **kwargs)
                request_meta["response_length"] = len(
                    request_meta["response"].data)
            except grpc.RpcError as e:
                request_meta["exception"] = e
            request_meta["response_time"] = (
                time.perf_counter() - start_perf_counter) * 1000
            self.env.events.request.fire(**request_meta)
            return request_meta["response"]

        return wrapper


class GrpcUser(User):
    abstract = True

    stub_class = None

    def __init__(self, environment):
        super().__init__(environment)
        for attr_value, attr_name in ((self.host, "host"), (self.stub_class, "stub_class")):
            if attr_value is None:
                raise LocustError(f"You must specify the {attr_name}.")
        self._channel = grpc.insecure_channel(self.host)
        self._channel_closed = False
        stub = self.stub_class(self._channel)
        self.client = GrpcClient(environment, stub)


class RequestGrpcUser(GrpcUser):
    host = "127.0.0.1:9898"
    stub_class = message_proto_pb2_grpc.EchoProtoServerStub

    @task
    def request(self):
        method = "send_order"
        if not self._channel_closed:
            res = self.client.request(
                message_proto_pb2.GateWayRequest(method=method))
            print("[RequestGrpcUser-res] ", res)


class PerformanceTest(unittest.TestCase):

    @classmethod
    def setUpClass(self):
        warnings.simplefilter('ignore', ResourceWarning)

    @unittest.skip
    def test_rpc(self):
        run()
        pass


# pip install locust
# 检验是否安装成功 locust -V
# 压力测试命令：locust -f grpc_client.py -c 4 -r 2  --run-time 15s --no-web
"""
web模式：
运行 locust -f grpc_client.py  填写user 10  hatch 2也就是每秒发多少请求
"""
if __name__ == "__main__":
    # unittest.main()
    # run()
    pass
